#!/usr/bin/env python

# Author:   Edward Roualdes
#           [2012.11.20]
#           University of Kentucky

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from collections import defaultdict
import utils


# function definition
def sortBarcode(INPUTFILE, BARCODEFILE, trim):
    """Sort DNA reads from INPUTFILE by individual's names in BARCODEFILE searching first for the barcode in the beginning 30 nucleotides and then, if previously unsuccesful, the reverse complement's first 30 nucleotides.

    input: 
    sorted            = list of tuples [(barcode_1, dnaRead_1), ... (barcode_n, dnaRead_n)]
    n                 = number of DNA reads
    nomatch           = number of reads with no matchin barcodes
    multmatch         = number of reads that matched multiple barcodes
    match             = number of reads that were matched

    output:
    c                 = number of reads sorted by barcode
    sortedBarcodes    = defaultdict( (barcode_1, dnaRead_1_1, ..., dnaRead_1_m), ...
    (barcode_c, dnaRead_c_1, ..., dnaRead_c_l) )
"""


    sortedBarcodes = defaultdict(list) # initialize the data structure to be returned

    # create list of barcode strings from BARCODEFILE
    # list parsing less memory efficient when we want all info per sequence
    bOpen = open(BARCODEFILE, "rU")
    bList = list(SeqIO.parse(bOpen, "fasta"))
    bOpen.close()
    # print len(bList)

    # creat list of dna reads from INPUTFILE
    iOpen = open(INPUTFILE, "rU")
    iList = list(SeqIO.parse(iOpen, "fasta"))
    iOpen.close()
    # print len(iList)

    # initializations
    n = len(iList)              # number of reads
    nomatch = 0
    multmatch = 0

    # match reads to barcodes
    for read in iList:          # for each read
        match = 0
        sequence = read.seq # temporary copy read
        l = len(sequence)   # length of read
        if l > 30:
            s = sequence.lower()[0:29] # first 30 nucleotides in read
            r = sequence.reverse_complement().lower()[0:29] # first 30 rev_comp
            for code in bList:
                bs = code.seq.lower() # barcode
                lb = len(bs)          # length barcode
                found = s.find(bs) # find index of bs in seq if present ow -1
                if found > -1:     # if barcode found within first 30
                    if trim:
                        # trim beg/end by found+lb
                        # read.seq = sequence[found+lb : l-found-lb]
                        # trim beg/end by 30
                        read.seq = sequence[30:l-30] # trim 30
                    else:
                        read.seq = sequence
                    match += 1
                    whichBarcode = code.name
                    break
                else:
                    found = r.find(bs) # find index of bs in rev_seq
                    if found > -1:
                        if trim:
                            # trim beg/end by found+lb
                            # read.seq = sequence[found+lb:l-found-lb]
                            # trim beg/end by 30
                            read.seq = sequence[30:l-30] # trim 30
                        else:
                            read.seq = sequence
                        match += 1
                        whichBarcode = code.name
                        break
            if match == 1:
                sortedBarcodes[whichBarcode].append(read)
            elif match == 0:
                nomatch += 1
            # else:
            #     multmatch += 1
        else:
            nomatch += 1


    c = n - nomatch - multmatch # no. BARCODE classified reads

    del iList, bList
    print "Classified", c, "reads by barcode..."
    return sortedBarcodes, c



# # test function: $ python sortBarcode.py [INPUTFILE] [BARCODEFILE]

# if __name__ == "__main__":
#     import sys
#     i = sys.argv[1]             # INPUTFILE
#     b = sys.argv[2]             # BARCODEFILE

#     # creat list of dna reads from INPUTFILE
#     iOpen = open(i, "rU")
#     iList = list(SeqIO.parse(iOpen, "fasta"))
#     iOpen.close()
    
#     srt, c = sortBarcode(i, b, True)
#     # print srt.items()[0][1][0]
#     # del srt.items()[0][1][0]
#     # print srt.items()[0][1][0]
#     print type(iList[0]), dir(iList[0])
        
#     # for inds, reads in srt.iteritems():
#     #     if c < limit:
#     #         for read in reads:
#     #             print read
            

#             # if read.id == 'HOBQI4X01AP0XM':
#             #     print read.id, '\n', read.seq

#     # list of tuples, each tuples' 1st element is barcode name
#     # for inds, reads in srt.iteritems():
#     #     for read in reads:
#     #         print read

#     # print the BarcodeNames
#     # print [x[0] for x in sortBarcode(i, b).items()]
#     # print map(lambda x: x[0], sortBarcode(i, b).items())
